# This script creates a connector for the DeepSeek model in OpenSearch Service.
# It uses the AWS4Auth library to sign the request with AWS credentials.

import boto3
import json
import os
import requests
from requests_aws4auth import AWS4Auth

opensearch_service_api_endpoint = os.environ['OPENSEARCH_SERVICE_DOMAIN_ENDPOINT']
region = os.environ['DEEPSEEK_AWS_REGION']
invoke_role_arn = os.environ['INVOKE_DEEPSEEK_ROLE']
create_deepseek_connector_role_arn = os.environ['CREATE_DEEPSEEK_CONNECTOR_ROLE']
sagemaker_endpoint_url = os.environ['SAGEMAKER_MODEL_INFERENCE_ENDPOINT']


# Create the AWS4Auth object that will sign the create connector API call. 
credentials = boto3.client('sts').assume_role(
  RoleArn=create_deepseek_connector_role_arn,
  RoleSessionName='create_connector_session'
)['Credentials']
awsauth = AWS4Auth(credentials['AccessKeyId'],
                   credentials['SecretAccessKey'],
                   region,
                   'es',
                   session_token=credentials['SessionToken'])


# Prepare the API call parameters. 
path = '/_plugins/_ml/connectors/_create'
url = opensearch_service_api_endpoint + path
# See the documentation 
# https://opensearch.org/docs/latest/ml-commons-plugin/remote-models/blueprints/ for 
# details on the connector payload, and additional blueprints for other models.
payload = {
  "name": "DeepSeek R1 model connector v2",
  "description": "Connector for my Sagemaker DeepSeek model",
  "version": "1.0",
  "protocol": "aws_sigv4",
  "credential": {
    "roleArn": invoke_role_arn
  },
  "parameters": {
    "service_name": "sagemaker",
    "region": region,
    "do_sample": True,
    "top_p": 0.7,
    "temperature": 0.5,
    "max_tokens": 1024
  },
  "actions": [
    {
      "action_type": "PREDICT",
      "method": "POST",
      "url": sagemaker_endpoint_url,
      "headers": {
        "content-type": "application/json"
      },
      "request_body": """{ "prompt": "${parameters.inputs}", "temperature": ${parameters.temperature}, "top_p": ${parameters.top_p}, "max_tokens": ${parameters.max_tokens} }""",
      "post_process_function": """
            try {
            // 从响应中提取生成的文本
            def completion = "";
            
            // 检查是否存在 choices 数组
            if (params.choices.size() > 0) {
              // 获取第一个选择的文本
              completion = params.choices[0].text;
              
            } else {
              completion = "无法从模型响应中提取文本";
            }
            
            
            def json = "{" +
                             "\\"name\\": \\"response\\","+
                            "\\"dataAsMap\\": {" +
                             "\\"completion\\":\\"" + escape(completion) + "\\"}" +
                             "}";
            return json;
            
        
          } catch (Exception e) {
            // 返回错误信息
            return '{' +
                     '"name": "response",' +
                     '"dataAsMap": {' +
                        '"completion": "处理错误: ' + escape(e) + '"' +
                     '}' +
                   '}';
          }
        """
    }
  ]
}

# This ignores errors and doesn't check the result. In real use,
# you should wrap this code with try/except blocks and check the
# response status code and the response body for errors.
headers = {"Content-Type": "application/json"}
r = requests.post(url, auth=awsauth, json=payload, headers=headers)
print(r.text)
connector_id = json.loads(r.text)['connector_id']


# print(connector_id)
print(f'DEEPSEEK_CONNECTOR_ID="{connector_id}"\n')
os.environ["DEEPSEEK_CONNECTOR_ID"] = connector_id